// 导入用于构建聊天完成请求的类
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
// 导入表示聊天消息的类
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
// 导入表示聊天消息角色的枚举类
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
// 导入Ark服务类，用于与Ark API进行交互
import com.volcengine.ark.runtime.service.ArkService;
// 导入OkHttp的连接池类，用于管理HTTP连接
import okhttp3.ConnectionPool;
// 导入OkHttp的调度器类，用于管理请求的执行
import okhttp3.Dispatcher;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * 这个类是一个示例，展示了如何使用VolcEngine Ark服务进行聊天完成请求，
 * 包括标准请求和流式请求两种方式。
 */
public class ChatCompletionsExample {
    // 从环境变量中获取Ark API密钥
    static String apiKey = System.getenv("ARK_API_KEY");
    // 创建一个连接池，最大连接数为5，空闲连接的存活时间为1秒
    static ConnectionPool connectionPool = new ConnectionPool(5, 1, TimeUnit.SECONDS);
    // 创建一个调度器，用于管理请求的执行
    static Dispatcher dispatcher = new Dispatcher();
    // 构建Ark服务实例，配置调度器、连接池、基础URL和API密钥
    static ArkService service = ArkService.builder()
            .dispatcher(dispatcher)
            .connectionPool(connectionPool)
            .baseUrl("https://ark.cn-beijing.volces.com/api/v3")
            .apiKey(apiKey)
            .build();

    public static void main(String[] args) {
        // 打印标准请求的提示信息
        System.out.println("\n----- standard request -----");
        // 创建一个用于存储聊天消息的列表
        final List<ChatMessage> messages = new ArrayList<>();
        // 创建一个系统消息，指定消息角色为系统，并设置消息内容
        final ChatMessage systemMessage = ChatMessage.builder()
                .role(ChatMessageRole.SYSTEM)
                .content("你是豆包，是由字节跳动开发的 AI 人工智能助手")
                .build();
        // 创建一个用户消息，指定消息角色为用户，并设置消息内容
        final ChatMessage userMessage = ChatMessage.builder()
                .role(ChatMessageRole.USER)
                .content("常见的十字花科植物有哪些？")
                .build();
        // 将系统消息添加到消息列表中
        messages.add(systemMessage);
        // 将用户消息添加到消息列表中
        messages.add(userMessage);

        // 构建聊天完成请求，指定使用的模型和消息列表
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model("ep-20250207210413-l6c9q")
                .messages(messages)
                .build();

        // 发送标准的聊天完成请求，遍历响应中的每个选择，并打印消息内容
        service.createChatCompletion(chatCompletionRequest)
                .getChoices()
                .forEach(choice -> System.out.println(choice.getMessage().getContent()));

        // 打印流式请求的提示信息
        System.out.println("\n----- streaming request -----");
        // 创建一个用于存储流式聊天消息的列表
        final List<ChatMessage> streamMessages = new ArrayList<>();
        // 创建一个流式系统消息，指定消息角色为系统，并设置消息内容
        final ChatMessage streamSystemMessage = ChatMessage.builder()
                .role(ChatMessageRole.SYSTEM)
                .content("你是豆包，是由字节跳动开发的 AI 人工智能助手")
                .build();
        // 创建一个流式用户消息，指定消息角色为用户，并设置消息内容
        final ChatMessage streamUserMessage = ChatMessage.builder()
                .role(ChatMessageRole.USER)
                .content("常见的十字花科植物有哪些？")
                .build();
        // 将流式系统消息添加到流式消息列表中
        streamMessages.add(streamSystemMessage);
        // 将流式用户消息添加到流式消息列表中
        streamMessages.add(streamUserMessage);

        // 构建流式聊天完成请求，指定使用的模型和消息列表
        ChatCompletionRequest streamChatCompletionRequest = ChatCompletionRequest.builder()
                .model("ep-20250207210413-l6c9q")
                .messages(streamMessages)
                .build();

        // 发送流式聊天完成请求，处理可能的错误，并逐块打印响应内容
        service.streamChatCompletion(streamChatCompletionRequest)
                .doOnError(Throwable::printStackTrace)
                .blockingForEach(
                        choice -> {
                            // 如果响应中的选择列表不为空
                            if (choice.getChoices().size() > 0) {
                                // 打印第一个选择的消息内容
                                System.out.print(choice.getChoices().get(0).getMessage().getContent());
                            }
                        }
                );

        // 所有请求完成后，关闭服务的执行器
        service.shutdownExecutor();
    }
}